import cv2
import numpy as np
import math

"""
{{a -> -((n sumxy - sumx sumy)/(sumx^2 - n sumx2)), 
  b -> -((-sumx sumxy + sumx2 sumy)/(sumx^2 - n sumx2))}}
"""
def leastSquaresLine(points):
    """最小二乘拟合直线"""
    sumx = 0
    sumy = 0
    sumxy = 0
    sumx2 = 0
    n = len(points)
    for i in range(n):
        x = points[i, 0, 0]
        y = points[i, 0, 1]
        sumx += x
        sumy += y
        sumx2 += x**2
        sumxy += x * y
    
    a = -((n * sumxy - sumx * sumy)/(sumx**2 - n * sumx2))
    b = (sumy - sumx * a) / n

    return (a, b)

def pointDistanceToLine(a,b,x,y):
    # y = ax + b ==>ax - y +b = 0
    # A = a B = -1 C = b
    v = abs(a * x - y + b)
    w = math.sqrt(a**2 + 1)
    return v / w


def filterPoint(a, b, points):
    result = []
    n = len(points)
    for i in range(n):
        x = points[i, 0, 0]
        y = points[i, 0, 1]
        if pointDistanceToLine(a, b,x , y) <= 4:
            result.append(i)

    if len(result) == n:
        return (False, [])
    else:
        newpoints = np.zeros((len(result), 1, 2), dtype=np.float)
        for i in range(len(result)):
            newpoints[i] = points[result[i]]

        return (True, newpoints)

def find_line(img_name):
    cx = 492
    cy = 575

    src = cv2.imread(img_name)
    src = cv2.cvtColor(src, cv2.COLOR_BGR2GRAY)
    roi_img = src[cy-30:cy+30, cx-200:cx+201]
    th1,binary_img1 = cv2.threshold(roi_img, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    kernel = np.ones((5,5),np.uint8)
    roi_img_closed = cv2.morphologyEx(binary_img1, cv2.MORPH_CLOSE, kernel, borderType=cv2.BORDER_REPLICATE)
    roi_img_opened = cv2.morphologyEx(roi_img_closed, cv2.MORPH_OPEN, kernel,borderType=cv2.BORDER_REPLICATE)
    roi_canny = cv2.Canny(roi_img_opened, 100, 180)
    lines = cv2.HoughLinesP(roi_canny,1,np.pi/360, 5,minLineLength=10,maxLineGap=5)

    # cloneTmage, contours, heriachy = cv2.findContours(roi_img_opened, cv2.RETR_LIST, cv2.CHAIN_APPROX_NONE)
    points = np.zeros((len(lines) * 2, 1, 2), dtype=np.float)
    for i in range(len(lines)):
        x1 = lines[i,0, 0]
        y1 = lines[i,0, 1]
        x2 = lines[i,0, 2]
        y2 = lines[i,0, 3]

        points[2*i, 0, 0] = x1
        points[2*i, 0, 1] = y1
        points[2*i+1, 0, 0] = x2
        points[2*i+1, 0, 1] = y2

    # for line in lines:
    #     x1 = line[0, 0]
    #     y1 = line[0, 1]
    #     x2 = line[0, 2]
    #     y2 = line[0, 3]
    #     cv2.line(roi_img, (x1, y1), (x2, y2), (255,0,0), 1)

    a = None
    b = None
    count = 0
    while True:
        count += 1
        a,b = leastSquaresLine(points)
        filtered, points = filterPoint(a, b, points)
        if filtered is False:
            break

        if count >=2:
            break
        
    print(a, b)
    p1 = (cx-200, cy-30 + int(b+0.5))
    p2 = (cx-200 + 800, cy-30 + int((800) * a + b + 0.5))

    cv2.line(src, p1, p2, (0,0,255), 1, 8)
    print(p1, p2)

    cv2.imshow("roi_img_opened", roi_img_opened)
    cv2.imshow("roi_img", roi_img)
    # cv2.imshow("roi_canny", roi_canny)
    cv2.imshow("src", src)

    cv2.waitKey()

if __name__ == "__main__":
    find_line("F:/new_code/450_copper_coil/opencv_find_line/Line_Image3.bmp")